#!/usr/bin/env python
# -*- coding: UTF-8 -*-
"""
@Project ：python_learning 
@File ：SQLServerORMGenerator.py
@IDE  ：PyCharm 
@Author ：李涵彬
@Date ：2025/7/9 下午3:21 
"""

import pyodbc
import os
import logging
from datetime import datetime
from typing import Dict, List, Tuple
import re

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# 数据库配置
SQLSERVER_DATABASE = 'ColdRollingFormal'
SQLSERVER_PASSWORD = '123456'
SQLSERVER_SERVER = '127.0.0.1'
SQLSERVER_PORT = '1433'
SQLSERVER_USER = 'sa'
SQLSERVER_CHARSET = 'UTF-8'


class SQLServerORMGenerator:
	"""SQL Server ORM代码生成器"""

	def __init__(self):
		self.connection = None
		self.output_dir = "models"
		self.base_class_name = "BaseModel"

		# SQL Server类型到Python类型的映射
		self.type_mapping = {
			'int': 'int',
			'bigint': 'int',
			'smallint': 'int',
			'tinyint': 'int',
			'bit': 'bool',
			'decimal': 'float',
			'numeric': 'float',
			'money': 'float',
			'smallmoney': 'float',
			'float': 'float',
			'real': 'float',
			'date': 'datetime.date',
			'time': 'datetime.time',
			'datetime': 'datetime.datetime',
			'datetime2': 'datetime.datetime',
			'smalldatetime': 'datetime.datetime',
			'timestamp': 'bytes',
			'binary': 'bytes',
			'varbinary': 'bytes',
			'image': 'bytes',
			'text': 'str',
			'ntext': 'str',
			'varchar': 'str',
			'nvarchar': 'str',
			'char': 'str',
			'nchar': 'str',
			'xml': 'str',
			'uniqueidentifier': 'str',
			'sql_variant': 'object'
		}

		# 需要导入的模块
		self.imports = {
			'datetime': ['datetime', 'date', 'time'],
			'typing': ['Optional', 'List'],
			'dataclasses': ['dataclass'],
			'decimal': ['Decimal']
		}

	def connect_to_database(self) -> bool:
		"""连接到SQL Server数据库"""
		try:
			# 构建连接字符串
			connection_string = (
				f"DRIVER={{ODBC Driver 17 for SQL Server}};"
				f"SERVER={SQLSERVER_SERVER},{SQLSERVER_PORT};"
				f"DATABASE={SQLSERVER_DATABASE};"
				f"UID={SQLSERVER_USER};"
				f"PWD={SQLSERVER_PASSWORD};"
				f"charset={SQLSERVER_CHARSET};"
			)

			self.connection = pyodbc.connect(connection_string)
			logger.info("数据库连接成功")
			return True

		except Exception as e:
			logger.error(f"数据库连接失败: {e}")
			return False

	def get_all_tables(self) -> List[str]:
		"""获取数据库中所有表名"""
		try:
			cursor = self.connection.cursor()
			query = """
            SELECT TABLE_NAME 
            FROM INFORMATION_SCHEMA.TABLES 
            WHERE TABLE_TYPE = 'BASE TABLE'
            ORDER BY TABLE_NAME
            """
			cursor.execute(query)
			tables = [row[0] for row in cursor.fetchall()]
			logger.info(f"找到 {len(tables)} 个表")
			return tables

		except Exception as e:
			logger.error(f"获取表列表失败: {e}")
			return []

	def get_table_schema(self, table_name: str) -> List[Dict]:
		"""获取表的schema信息"""
		try:
			cursor = self.connection.cursor()
			query = """
            SELECT 
                COLUMN_NAME,
                DATA_TYPE,
                IS_NULLABLE,
                COLUMN_DEFAULT,
                CHARACTER_MAXIMUM_LENGTH,
                NUMERIC_PRECISION,
                NUMERIC_SCALE,
                ORDINAL_POSITION
            FROM INFORMATION_SCHEMA.COLUMNS
            WHERE TABLE_NAME = ?
            ORDER BY ORDINAL_POSITION
            """
			cursor.execute(query, (table_name,))

			columns = []
			for row in cursor.fetchall():
				column_info = {
					'name': row[0],
					'type': row[1],
					'nullable': row[2] == 'YES',
					'default': row[3],
					'max_length': row[4],
					'precision': row[5],
					'scale': row[6],
					'position': row[7]
				}
				columns.append(column_info)

			return columns

		except Exception as e:
			logger.error(f"获取表 {table_name} 的schema失败: {e}")
			return []

	def get_primary_keys(self, table_name: str) -> List[str]:
		"""获取表的主键列"""
		try:
			cursor = self.connection.cursor()
			query = """
            SELECT c.COLUMN_NAME
            FROM INFORMATION_SCHEMA.TABLE_CONSTRAINTS tc
            JOIN INFORMATION_SCHEMA.CONSTRAINT_COLUMN_USAGE ccu ON tc.CONSTRAINT_NAME = ccu.CONSTRAINT_NAME
            JOIN INFORMATION_SCHEMA.COLUMNS c ON ccu.COLUMN_NAME = c.COLUMN_NAME AND ccu.TABLE_NAME = c.TABLE_NAME
            WHERE tc.TABLE_NAME = ? AND tc.CONSTRAINT_TYPE = 'PRIMARY KEY'
            """
			cursor.execute(query, (table_name,))
			return [row[0] for row in cursor.fetchall()]

		except Exception as e:
			logger.error(f"获取表 {table_name} 的主键失败: {e}")
			return []

	def get_foreign_keys(self, table_name: str) -> List[Dict]:
		"""获取表的外键信息"""
		try:
			cursor = self.connection.cursor()
			query = """
	        SELECT 
	            fk.name AS constraint_name,
	            tp.name AS parent_table,
	            cp.name AS parent_column,
	            tr.name AS referenced_table,
	            cr.name AS referenced_column
	        FROM sys.foreign_keys fk
	        INNER JOIN sys.foreign_key_columns fkc ON fk.object_id = fkc.constraint_object_id
	        INNER JOIN sys.tables tp ON fkc.parent_object_id = tp.object_id
	        INNER JOIN sys.columns cp ON fkc.parent_object_id = cp.object_id AND fkc.parent_column_id = cp.column_id
	        INNER JOIN sys.tables tr ON fkc.referenced_object_id = tr.object_id
	        INNER JOIN sys.columns cr ON fkc.referenced_object_id = cr.object_id AND fkc.referenced_column_id = cr.column_id
	        WHERE tp.name = ?
	        """
			cursor.execute(query, (table_name,))

			foreign_keys = []
			for row in cursor.fetchall():
				fk_info = {
					'constraint_name': row[0],
					'column': row[2],
					'referenced_table': row[3],
					'referenced_column': row[4]
				}
				foreign_keys.append(fk_info)

			return foreign_keys

		except Exception as e:
			logger.error(f"获取表 {table_name} 的外键失败: {e}")
			return []

	def convert_type(self, sql_type: str, nullable: bool = True) -> str:
		"""将SQL Server类型转换为Python类型"""
		python_type = self.type_mapping.get(sql_type.lower(), 'str')

		if nullable and python_type != 'object':
			return f"Optional[{python_type}]"
		return python_type

	def snake_to_pascal(self, snake_str: str) -> str:
		"""将snake_case转换为PascalCase"""
		words = snake_str.split('_')
		return ''.join(word.capitalize() for word in words)

	def generate_model_class(self, table_name: str) -> str:
		"""生成单个表的ORM类代码"""
		# 获取表信息
		columns = self.get_table_schema(table_name)
		primary_keys = self.get_primary_keys(table_name)
		foreign_keys = self.get_foreign_keys(table_name)

		if not columns:
			logger.warning(f"表 {table_name} 没有列信息")
			return ""

		# 类名
		class_name = self.snake_to_pascal(table_name)

		# 生成代码
		code_lines = []

		# 文件头部注释
		code_lines.append('"""')
		code_lines.append(f'自动生成的ORM模型类 - {table_name}')
		code_lines.append(f'生成时间: {datetime.now().strftime("%Y-%m-%d %H:%M:%S")}')
		code_lines.append('"""')
		code_lines.append('')

		# 导入语句
		used_imports = set()
		for column in columns:
			python_type = self.convert_type(column['type'], column['nullable'])
			if 'datetime' in python_type:
				used_imports.add('datetime')
			if 'Optional' in python_type:
				used_imports.add('typing')
			if 'Decimal' in python_type:
				used_imports.add('decimal')

		if 'datetime' in used_imports:
			code_lines.append('from datetime import datetime, date, time')
		if 'typing' in used_imports:
			code_lines.append('from typing import Optional, List')
		if 'decimal' in used_imports:
			code_lines.append('from decimal import Decimal')

		code_lines.append('from dataclasses import dataclass')
		code_lines.append('')

		# 类定义
		code_lines.append('@dataclass')
		code_lines.append(f'class {class_name}:')
		code_lines.append(f'    """')
		code_lines.append(f'    {table_name}表的ORM模型类')
		code_lines.append(f'    """')
		code_lines.append('')

		# 表信息
		code_lines.append(f'    # 表名')
		code_lines.append(f'    __tablename__ = "{table_name}"')
		code_lines.append('')

		# 主键信息
		if primary_keys:
			code_lines.append(f'    # 主键: {", ".join(primary_keys)}')

		# 外键信息
		if foreign_keys:
			code_lines.append(f'    # 外键关系:')
			for fk in foreign_keys:
				code_lines.append(f'    # {fk["column"]} -> {fk["referenced_table"]}.{fk["referenced_column"]}')

		code_lines.append('')

		# 字段定义
		for column in columns:
			comment_parts = []

			# 类型信息
			type_info = column['type']
			if column['max_length']:
				type_info += f"({column['max_length']})"
			elif column['precision'] and column['scale']:
				type_info += f"({column['precision']},{column['scale']})"

			comment_parts.append(f"SQL类型: {type_info}")

			# 是否可空
			if not column['nullable']:
				comment_parts.append("NOT NULL")

			# 默认值
			if column['default']:
				comment_parts.append(f"默认值: {column['default']}")

			# 主键标识
			if column['name'] in primary_keys:
				comment_parts.append("主键")

			# 外键标识
			fk_refs = [fk for fk in foreign_keys if fk['column'] == column['name']]
			if fk_refs:
				fk_ref = fk_refs[0]
				comment_parts.append(f"外键 -> {fk_ref['referenced_table']}.{fk_ref['referenced_column']}")

			comment = " | ".join(comment_parts)
			python_type = self.convert_type(column['type'], column['nullable'])

			code_lines.append(f'    {column["name"]}: {python_type}  # {comment}')

		code_lines.append('')

		# 生成一些实用方法
		code_lines.append('    def to_dict(self) -> dict:')
		code_lines.append('        """转换为字典"""')
		code_lines.append('        return {')
		for column in columns:
			code_lines.append(f'            "{column["name"]}": self.{column["name"]},')
		code_lines.append('        }')
		code_lines.append('')

		code_lines.append('    @classmethod')
		code_lines.append('    def from_dict(cls, data: dict):')
		code_lines.append('        """从字典创建对象"""')
		code_lines.append('        return cls(')
		for column in columns:
			code_lines.append(f'            {column["name"]}=data.get("{column["name"]}"),')
		code_lines.append('        )')
		code_lines.append('')

		code_lines.append('    def __repr__(self) -> str:')
		code_lines.append('        """字符串表示"""')
		if primary_keys:
			pk_repr = ', '.join([f'{pk}={{{pk}}}' for pk in primary_keys])
			code_lines.append(f'        return f"{class_name}({pk_repr})"')
		else:
			code_lines.append(f'        return f"{class_name}({{id(self)}})"')

		return '\n'.join(code_lines)

	def generate_init_file(self, table_names: List[str]) -> str:
		"""生成__init__.py文件"""
		code_lines = []

		# 文件头部注释
		code_lines.append('"""')
		code_lines.append('自动生成的模型包初始化文件')
		code_lines.append(f'生成时间: {datetime.now().strftime("%Y-%m-%d %H:%M:%S")}')
		code_lines.append('"""')
		code_lines.append('')

		# 导入所有模型
		for table_name in sorted(table_names):
			class_name = self.snake_to_pascal(table_name)
			code_lines.append(f'from .{table_name} import {class_name}')

		code_lines.append('')

		# __all__ 定义
		code_lines.append('__all__ = [')
		for table_name in sorted(table_names):
			class_name = self.snake_to_pascal(table_name)
			code_lines.append(f'    "{class_name}",')
		code_lines.append(']')

		return '\n'.join(code_lines)

	def generate_models(self):
		"""生成所有表的ORM模型"""
		if not self.connect_to_database():
			return

		try:
			# 创建输出目录
			if not os.path.exists(self.output_dir):
				os.makedirs(self.output_dir)
				logger.info(f"创建输出目录: {self.output_dir}")

			# 获取所有表
			tables = self.get_all_tables()
			if not tables:
				logger.warning("没有找到任何表")
				return

			successful_tables = []
			failed_tables = []

			# 为每个表生成ORM类
			for table_name in tables:
				try:
					logger.info(f"正在生成 {table_name} 的ORM类...")

					# 生成模型代码
					model_code = self.generate_model_class(table_name)

					if model_code:
						# 写入文件
						file_path = os.path.join(self.output_dir, f"{table_name}.py")
						with open(file_path, 'w', encoding='utf-8') as f:
							f.write(model_code)

						successful_tables.append(table_name)
						logger.info(f"✓ {table_name} 生成成功")
					else:
						failed_tables.append(table_name)
						logger.error(f"✗ {table_name} 生成失败")

				except Exception as e:
					logger.error(f"生成 {table_name} 时发生错误: {e}")
					failed_tables.append(table_name)

			# 生成__init__.py文件
			try:
				init_code = self.generate_init_file(successful_tables)
				init_path = os.path.join(self.output_dir, "__init__.py")
				with open(init_path, 'w', encoding='utf-8') as f:
					f.write(init_code)
				logger.info("✓ __init__.py 生成成功")
			except Exception as e:
				logger.error(f"生成__init__.py失败: {e}")

			# 输出结果摘要
			logger.info("=" * 60)
			logger.info("生成结果摘要:")
			logger.info(f"成功生成: {len(successful_tables)} 个模型")
			logger.info(f"失败: {len(failed_tables)} 个模型")

			if successful_tables:
				logger.info("成功生成的模型:")
				for table in successful_tables:
					logger.info(f"  - {table}")

			if failed_tables:
				logger.info("失败的模型:")
				for table in failed_tables:
					logger.info(f"  - {table}")

		except Exception as e:
			logger.error(f"生成过程中发生错误: {e}")

		finally:
			if self.connection:
				self.connection.close()
				logger.info("数据库连接已关闭")


def main():
	"""主函数"""
	print("SQL Server ORM模型生成器")
	print("=" * 50)

	generator = SQLServerORMGenerator()
	generator.generate_models()


if __name__ == "__main__":
	main()
